package SimpleLACore

import chisel3._
import chisel3.util._
import decode._
import ALU._
import Consts._

class DebugPort extends Bundle{
  val pc = Output(UInt(32.W))
  val wen = Output(UInt(4.W))
  val wnum = Output(UInt(5.W))
  val wdata = Output(UInt(32.W))
  val inst = Output(UInt(32.W))
}

class SimpleLACore(useDiff:Boolean = true) extends Module with SimpleLACoreParam {
  val io = IO(new Bundle{
    val ipi = Input(Bool())
    val interrupt = Input(UInt(8.W))
    val inst = Flipped(new SimpleICacheIO)
    val data = Flipped(new SimpleCacheIO)
    val debug = Output(new DebugPort)
    val diff = if(useDiff) Output(new Difftest) else null
  })

  val idle = RegInit(false.B)
  val PC = RegInit(initPC)
  val PC4 = PC + 4.U
  val GR = Mem(32, UInt(32.W))
//  val rand = random.LFSR(16)
  val timer = RegInit(0.U(64.W))
  timer := timer + 1.U

  val iStall = WireInit(false.B)
  val dStall = WireInit(false.B)
  val iStallReg = RegNext(iStall, false.B)
  val dStallReg = RegNext(dStall, false.B)


  //CSRs
  val crmd = RegInit(8.U.asTypeOf(new CRMD))
  val prmd = RegInit(0.U.asTypeOf(new PRMD))
  val euen = RegInit(0.U.asTypeOf(new EUEN))
  val ectl = RegInit(0.U.asTypeOf(new ECTL))
  val estat = RegInit(0.U.asTypeOf(new ESTAT))
  val era = RegInit(0.U.asTypeOf(new ERA))
  val badv = RegInit(0.U.asTypeOf(new BADV))
  val eentry = RegInit(0.U.asTypeOf(new EENTRY))
  val tlbidx = RegInit(0.U.asTypeOf(new TLBIDX))
  val tlbehi = RegInit(0.U.asTypeOf(new TLBEHI))
  val tlbelo0 = RegInit(0.U.asTypeOf(new TLBELO0))
  val tlbelo1 = RegInit(0.U.asTypeOf(new TLBELO1))
  val asid = RegInit(0x2800.U.asTypeOf(new ASID))
  val pgdl = RegInit(0.U.asTypeOf(new PGDL))
  val pgdh = RegInit(0.U.asTypeOf(new PGDH))
  val pgd = RegInit(0.U.asTypeOf(new PGD))
  val cpuid = RegInit(0.U.asTypeOf(new CPUID))
  val save0 = RegInit(0.U.asTypeOf(new SAVE0))
  val save1 = RegInit(0.U.asTypeOf(new SAVE1))
  val save2 = RegInit(0.U.asTypeOf(new SAVE2))
  val save3 = RegInit(0.U.asTypeOf(new SAVE3))
  val tid = RegInit(0.U.asTypeOf(new TID))
  val tcfg = RegInit(0.U.asTypeOf(new TCFG))
  val tval = RegInit(0.U.asTypeOf(new TVAL))
  val ticlr = RegInit(0.U.asTypeOf(new TICLR))
  val llbctl = RegInit(0.U.asTypeOf(new LLBCTL))
  val tlbrentry = RegInit(0.U.asTypeOf(new TLBRENTRY))
  val dmw0 = RegInit(0.U.asTypeOf(new DMW0))
  val dmw1 = RegInit(0.U.asTypeOf(new DMW1))

  val csrs = Seq(crmd, prmd, euen, ectl, estat, era, badv, eentry, tlbidx, tlbehi, tlbelo0, tlbelo1, asid,
    pgdl, pgdh, pgd, cpuid, save0, save1, save2, save3, tid, tcfg, tval, ticlr, llbctl, tlbrentry, dmw0, dmw1)

  implicit val csrFile = new CSRFile(csrs:_*)

  estat.IS(12) := io.ipi
  (estat.IS.tail.tail zip io.interrupt.asBools).foreach{case (is, int) => is := int}

  //Timer
  when(tcfg.En){
    when(tval.TimeVal =/= 0.U){
      tval.TimeVal := tval.TimeVal - 1.U
      when(tval.TimeVal === 1.U){
        estat.IS(11) := true.B
      }
    }.elsewhen(tcfg.Periodic) {
      tval.TimeVal := Cat(tcfg.InitVal, 0.U(2.W))
    }
  }

  //exceptions
  val INT = WireInit(false.B)
  val PIL = WireInit(false.B)
  val PIS = WireInit(false.B)
  val PIF = WireInit(false.B)
  val PME = WireInit(false.B)
  val PPI = WireInit(false.B)
  val ADEF = WireInit(false.B)
  val ADEM = WireInit(false.B)
  val ALE = WireInit(false.B)
  val SYS = WireInit(false.B)
  val BRK = WireInit(false.B)
  val INE = WireInit(false.B)
  val IPE = WireInit(false.B)
  val FPD = WireInit(false.B)
  val FPE = WireInit(false.B)
  val TLBR = WireInit(false.B)

  val excp = Cat(INT, PIL, PIS, PIF, PME, PPI, ADEF, ADEM, ALE, SYS, BRK, INE, IPE, FPD, FPE, TLBR).orR
  INT := (ectl.LIE & Cat(estat.IS.reverse)).orR && crmd.IE

  //TLB
//  val tlb = RegInit(VecInit(Seq.fill(nTLB)(0.U.asTypeOf(new TLB))))
  val tlb = Reg(Vec(nTLB,new TLB))

  def addr_trans(va: UInt, asid: UInt) = {
    val pa = Wire(UInt(32.W))
    val mat = Wire(UInt(2.W))
    val ade = WireInit(false.B)
    val miss = WireInit(false.B)
    val invalid = WireInit(false.B)
    val ppi = WireInit(false.B)
    val pme = WireInit(false.B)
    val asidMatch = tlb.map(_.ASID === asid)
    val vaMatch = tlb.map(i => i.VPPN(18,10) === va(31, 23) && (i.PS === 22.U || i.VPPN(9,0) === va(22, 13)))
    val tlbHit = (tlb zip asidMatch zip vaMatch).map{case ((t, a), v) => t.E && v && (t.G || a)}

    when(va(31, 29) === dmw0.VSEG && (dmw0.PLV0 && crmd.PLV === 0.U || dmw0.PLV3 && crmd.PLV === 3.U)){
      pa := Cat(dmw0.PSEG, va(28, 0))
      mat := dmw0.MAT
    }.elsewhen(va(31, 29) === dmw1.VSEG && (dmw1.PLV0 && crmd.PLV === 0.U || dmw1.PLV3 && crmd.PLV === 3.U)){
      pa := Cat(dmw1.PSEG, va(28, 0))
      mat := dmw1.MAT
    }.otherwise{
      val foundTLB = Mux1H(tlbHit, tlb)
      val oddPG = Mux(foundTLB.PS === 12.U, va(12), va(22))
      val foundP = Mux(oddPG, foundTLB.P1, foundTLB.P0)
      pa := Mux(foundTLB.PS ===12.U, Cat(foundP.PPN, va(11, 0)), Cat(foundP.PPN(19, 10), va(21, 10)))
      mat := foundP.MAT
      ade := false.B//crmd.PLV === 3.U && va(31)
      miss := !tlbHit.reduce(_|_)
      invalid := !foundP.V
      ppi := crmd.PLV > foundP.PLV
      pme := !foundP.D
    }
    (pa, mat, ade, miss, invalid, ppi, pme, asidMatch, vaMatch, tlbHit)
  }



//  IF
  val PCTrans = addr_trans(PC, asid.ASID)
  val IF_OK = WireInit(false.B)
  io.inst.req.valid := !iStallReg && !dStallReg && IF_OK
  io.inst.req.bits := DontCare
//  io.inst.req.bits.wen := 0.U(4.W)
  io.inst.req.bits.mat := Mux(crmd.DA, crmd.DATF, PCTrans._2)
  io.inst.req.bits.addr := Mux(crmd.DA, PC, PCTrans._1)
//  io.inst.req.bits.wdata := DontCare
//  io.inst.req.bits.cacop := false.B
//  io.inst.req.bits.code := DontCare
//  io.inst.req.bits.preld := false.B
  when(dStallReg || iStallReg){
    IF_OK := true.B
  }.elsewhen(!INT && !idle){
    when(PC(1, 0) =/= 0.U){
      ADEF := true.B
    }.elsewhen(crmd.DA){
      IF_OK := true.B
    }.elsewhen(PCTrans._3){
      ADEF := true.B
    }.elsewhen(PCTrans._4){
      TLBR := true.B
    }.elsewhen(PCTrans._5){
      PIF := true.B
    }.elsewhen(PCTrans._6){
      PPI := true.B
    }.otherwise{
      IF_OK := true.B
    }
    when(!IF_OK){
      badv.VAddr := PC
      when(!ADEF){
        tlbehi.VPPN := PC(31, 13)
      }
    }
  }

  iStall := !io.inst.resp.valid && (io.inst.req.valid || iStallReg)

//  ID
  val inst = Wire(UInt(32.W))
  val inst_reg = RegNext(inst)
  inst := Mux(dStallReg, inst_reg, io.inst.resp.bits)
  val UI12 = inst(21, 10)
  val SI12 = Cat(Fill(20, inst(21)),inst(21, 10))
  val SI14 = Cat(Fill(18, inst(23)),inst(23, 10))
  val OFF16 = Cat(Fill(14, inst(25)),inst(25, 10), 0.U(2.W))
  val OFF26 = Cat(Fill(4, inst(9)), inst(9, 0), inst(25, 10), 0.U(2.W))
  val UI20 = Cat(inst(24, 5), 0.U(12.W))
  val csrAddr = inst(23, 10)
  val j = inst(9, 5)
  val k = inst(14, 10)
  val d = inst(4, 0)
  //Use freechips.rocketchip.rocket.DecodeLogic
  val (instV: Bool) :: (instP: Bool) :: op1Sel :: op2Sel :: func :: c0 = ListLookup(inst, invalidInst, decodeTable)
  val wbSel :: wbDest :: memOp :: memType :: (isCSR: Bool) :: (break: Bool) :: (sysCall: Bool) :: c1 = c0
  val tlbOp :: (isERTN: Bool) :: (isIdle: Bool) :: brType :: (isCACOP: Bool) :: (isPRELD: Bool) :: Nil = c1
  val rj = Mux(j === 0.U, 0.U, GR.read(j))
  val kd = Mux(!memType(2,1).andR | isCSR | brType.orR, d, k)
  val rkd = Mux(kd === 0.U, 0.U, GR.read(kd))
  val inv_op_decode = Seq(d === 0.U || d === 1.U, d === 2.U, d === 3.U, d === 4.U, d === 5.U, d === 6.U)
  val ID_OK = WireInit(false.B)
  when(dStallReg){
    ID_OK := true.B
  }.elsewhen(IF_OK && !iStall){
    SYS := sysCall
    BRK := break
    INE := !instV || (tlbOp === tlb_inv && !inv_op_decode.reduce(_|_))
    if(IPEEnable) IPE := (instP || isCACOP && d =/= 8.U && d =/= 9.U) && crmd.PLV =/= 0.U
    when(!SYS && !BRK && !INE && !IPE){
      ID_OK := true.B
    }
  }

  //alu
  val aluOp1 = Mux(op1Sel === 0.U, rj, PC)
  val aluOp2 = VecInit(rkd, k, UI12, SI12, SI14, OFF16, UI20, OFF26)(op2Sel)
  val aluOut = alu(aluOp1, aluOp2, func)

  val EXASID = Mux(tlbOp === tlb_inv, rj(9, 0), asid.ASID)
  val EXVA = Mux1H(Seq(
    (memType =/= MT_X) -> aluOut,
    (tlbOp === tlb_srch) -> Cat(tlbehi.VPPN, 0.U(13.W)),
    (tlbOp === tlb_inv) -> rkd
  ))
  val EXTrans = addr_trans(EXVA, EXASID)

  //mem
  val mem_OK = WireInit(false.B)
  val memALE = aluOut(0) && memType(2, 1) === 1.U || aluOut(1, 0).orR && memType(2, 1) === 2.U
  val preldNop = WireInit(isPRELD)//Current inst is invalid preld so it becomes Nop
  val memMask = VecInit(1.U(4.W), 3.U(4.W), 0xf.U(4.W), 0.U(4.W))(memType(2, 1)) << aluOut(1, 0)

  io.data.req.valid := !dStallReg && mem_OK && memType =/= MT_X && !preldNop && !(isCACOP && d(2,0) =/= 1.U)
  io.data.req.bits.wen := Mux(memOp === M_ST || (memOp === M_SC && llbctl.ROLLB), memMask, 0.U(4.W))
  io.data.req.bits.mat := Mux(crmd.PG, EXTrans._2, crmd.DATM)
  io.data.req.bits.addr := Mux(isCACOP && d(4,3) === 2.U || !isCACOP && crmd.PG, EXTrans._1, aluOut)
  io.data.req.bits.wdata := rkd << (aluOut(1, 0) << 3)
  io.data.req.bits.cacop := isCACOP
  io.data.req.bits.code := d(4,3)
  io.data.req.bits.preld := isPRELD

  io.inst.cacop.valid := !dStallReg && mem_OK && isCACOP && d(2,0) === 0.U
  io.inst.cacop.bits := DontCare
  io.inst.cacop.bits.code := d(4,3)
  io.inst.cacop.bits.addr := Mux(isCACOP && d(4,3) === 2.U || !isCACOP && crmd.PG, EXTrans._1, aluOut)
  val shiftData = io.data.resp.bits >> (aluOut(1, 0) << 3)
  val extendData = Mux1H(Seq(
    (memType === MT_B ) -> Cat(Fill(24, shiftData(7)), shiftData(7, 0)),
    (memType === MT_BU) -> Cat(0.U(24.W), shiftData(7, 0)),
    (memType === MT_H ) -> Cat(Fill(16, shiftData(15)), shiftData(15, 0)),
    (memType === MT_HU) -> Cat(0.U(16.W), shiftData(15, 0)),
    (memType === MT_W ) -> shiftData
  ))
//  val memOut = Mux(memOp === M_SC, llbctl.ROLLB,  extendData)
  val memOut = extendData

  when(dStallReg){
    mem_OK := true.B
  }.elsewhen(ID_OK){
    when(isPRELD){
      mem_OK := true.B
    }
    when(memType === MT_X || isCACOP && d(4,3) =/= 2.U){
      mem_OK := true.B
    }.elsewhen(memALE){
      ALE := !isPRELD
    }.elsewhen(crmd.DA){
      mem_OK := true.B
    }.elsewhen(EXTrans._3){
      ADEM := !isPRELD
    }.elsewhen(EXTrans._4){
      TLBR := !isPRELD
    }.elsewhen(EXTrans._5){
      when(memOp(0)){//memOp === M_LD || memOp === M_LL
        PIS := !isPRELD
      }.otherwise{
        PIL := !isPRELD
      }
    }.elsewhen(EXTrans._6){
      PPI := !isPRELD
    }.elsewhen(EXTrans._7 && memOp(0)){
      PME := !isPRELD
    }.otherwise{
      mem_OK := true.B
      preldNop := false.B
    }
    when(!mem_OK){
      badv.VAddr := aluOut
      when(!ALE && !ADEM){
        tlbehi.VPPN := aluOut(31, 13)
      }
    }
  }
  when(memOp === M_LL && mem_OK && !dStall){
    llbctl.ROLLB := true.B
  }
  when(memOp === M_SC && mem_OK && !dStall){
    llbctl.ROLLB := false.B
  }
  dStall := !io.data.resp.valid && (io.data.req.valid || dStallReg)

  //CSR
  val csrRD = csrFile.read(csrAddr)
  //CSR inst don't trigger dStall
  val csrMask = Mux(j === 1.U, -1.S(32.W).asUInt, rj)
  when(ID_OK && isCSR & csrMask.orR){
    csrFile.write(csrAddr, rkd, csrMask)
  }

  //TLB
  //TLB inst don't trigger dStall
  when(ID_OK){
    switch(tlbOp){
      is(tlb_srch){
        when(EXTrans._10.reduce(_||_)){
          tlbidx.Index := Mux1H(EXTrans._10.zipWithIndex.map{ case (bool, i) => (bool, i.U(lgnTLB.W))})
          tlbidx.NE := false.B
        }.otherwise{
          tlbidx.NE := true.B
        }
      }
      is(tlb_rd){
        tlb(tlbidx.Index).read()
      }
      is(tlb_wr){
        tlb(tlbidx.Index).write()
      }
      is(tlb_fill){
//        tlb(rand).write()
        tlb(timer(lgnTLB - 1, 0)).write()
      }
      is(tlb_inv){
        (tlb zip EXTrans._8 zip EXTrans._9).foreach({ case ((t, a), v) =>
          t.E := t.E & ! Mux1H(inv_op_decode, Seq(
            true.B,
            t.G,
            !t.G,
            !t.G && a,
            !t.G && a && v,
            (!t.G || a) && v
          ))
        })
      }
    }
  }

  //WB
  io.debug.pc := PC
  io.debug.wen := 0.U(4.W)
  io.debug.wnum := DontCare
  io.debug.wdata := DontCare
  io.debug.inst := inst
  val wbIdx = Mux1H(Seq(
    (wbDest === WB_R0) -> 0.U(5.W),
    (wbDest === WB_R1) -> 1.U(5.W),
    (wbDest === WB_RD) -> d,
    (wbDest === WB_RJ) -> j,
  ))
  val wbData = Mux1H(Seq(
    (wbSel === WB_ALU  ) -> aluOut,
    (wbSel === WB_MEM  ) -> memOut,
    (wbSel === WB_PC4  ) -> PC4,
    (wbSel === WB_CSR  ) -> csrRD,
    (wbSel === WB_CNTID) -> tid.TID,
    (wbSel === WB_CNTVH) -> timer(63,32),
    (wbSel === WB_CNTVL) -> timer(31,0),
    (wbSel === WB_LLB  ) -> llbctl.ROLLB,
  ))
  when(mem_OK && wbIdx.orR && !dStall){
    GR.write(wbIdx, wbData)
    io.debug.wen := 0xf.U(4.W)
    io.debug.wnum := wbIdx
    io.debug.wdata := wbData
  }

  //ERTN
  when(ID_OK && isERTN){
    crmd.PLV := prmd.PPLV
    crmd.IE := prmd.PIE
    when(!llbctl.KLO){
      llbctl.ROLLB := false.B
    }
    llbctl.KLO := false.B
    when(estat.Ecode === 0x3f.U){
      crmd.PG := true.B
      crmd.DA := false.B
    }
  }

  when(ID_OK && isIdle){
    idle := true.B
  }

  //deal exceptions
  val ecodeNext = Mux1H(Seq(
    INT -> 0x0.U,
    ADEF -> 0x8.U,
    PIF -> 0x3.U,
    INE -> 0xd.U,
    IPE -> 0xe.U,
    SYS -> 0xB.U,
    BRK -> 0xc.U,
    ALE -> 0x9.U,
    ADEM -> 0x8.U,
    PIL -> 0x1.U,
    PIS -> 0x2.U,
    PME -> 0x4.U,
    PPI -> 0x7.U,
    FPD -> 0xf.U,
    FPE -> 0x12.U,
    TLBR -> 0x3F.U
  ))
  when(!iStall && !dStall) {
    when(excp) {
      idle := false.B
      prmd.PPLV := crmd.PLV
      prmd.PIE := crmd.IE
      crmd.PLV := 0.U
      crmd.IE := false.B
      era.PC := PC
      estat.Ecode := ecodeNext
      estat.EsubCode := Mux1H(Seq(
        ADEF -> 0.U,
        ADEM -> 1.U,
        FPE -> 0.U
      ))
    }
    when(TLBR) {
      crmd.DA := true.B
      crmd.PG := false.B
    }
  }

  //nPC
  when(!iStall && !dStall){
    when(TLBR){
      PC := tlbrentry.value()
    }.elsewhen(excp){
      PC := eentry.value()
    }.elsewhen(isERTN){
      PC := era.PC
    }.elsewhen(!idle/* && !(isIdle && ID_OK)*/){//when idle, mark exception on idle?
      val take = VecInit(
        false.B,
        rj =/= rkd,
        rj === rkd,
        rj.asSInt >= rkd.asSInt,
        rj >= rkd,
        rj.asSInt < rkd.asSInt,
        rj < rkd,
        true.B
      )(brType)
      PC := Mux(take, aluOut, PC4)
    }
  }

  when(reset.asBool){
    io.inst.req.valid := false.B
    io.data.req.valid := false.B
  }

  if(useDiff){
    val diff_inst  = Reg(new DifftestInstrCommit())
    val diff_excp  = Reg(new DifftestExcpEvent())
    val diff_store = Reg(new DifftestStoreEvent())
    val diff_load  = Reg(new DifftestLoadEvent())

    diff_inst.valid := mem_OK && !iStall && !dStall && !reset.asBool
    diff_inst.pc := PC
    diff_inst.instr := inst
    diff_inst.is_TLBFILL := tlbOp === tlb_fill
    diff_inst.TLBFILL_index := timer(lgnTLB - 1, 0)
    diff_inst.is_CNTinst := wbSel === WB_CNTID || wbSel === WB_CNTVH || wbSel === WB_CNTVL
    diff_inst.timer_64_value := timer
    diff_inst.wen := mem_OK && wbIdx.orR && !dStall
    diff_inst.wdest := wbIdx
    diff_inst.wdata := wbData
    diff_inst.csr_rstat := isCSR && csrAddr === estat.addr.U
    diff_inst.csr_data := csrRD
    diff_excp.excp_valid := excp && !iStall && !dStall && !reset.asBool
    diff_excp.eret := isERTN
    diff_excp.intrNo := estat.value()(csrFile)(12,2)
    diff_excp.cause := ecodeNext
    diff_excp.exceptionPC := PC
    diff_excp.exceptionInst := inst
    diff_store.valid := Cat(0.U(4.W),
      llbctl.ROLLB && memOp === M_SC,
      Cat(memType === MT_W,
        memType === MT_H,
        memType === MT_B) & Fill(3, memOp === M_ST)) &
      Fill(8, mem_OK && !iStall && !dStall && !reset.asBool)
    diff_store.storePAddr := io.data.req.bits.addr
    diff_store.storeVAddr := EXVA
    diff_store.storeData := io.data.req.bits.wdata & Cat(io.data.req.bits.wen.asBools.map(Fill(8,_)).reverse)
    diff_load.valid := Cat(0.U(2.W),
      memOp === M_LL,
      Cat(memType === MT_W,
        memType === MT_HU,
        memType === MT_H,
        memType === MT_BU,
        memType === MT_B) & Fill(5, memOp === M_LD)) &
      Fill(8, mem_OK && !iStall && !dStall && !reset.asBool)
    diff_load.paddr := io.data.req.bits.addr
    diff_load.vaddr := EXVA

    io.diff.inst := diff_inst
    io.diff.excp := diff_excp
    io.diff.store := diff_store
    io.diff.load := diff_load

    io.diff.csr.crmd := crmd.value()
    io.diff.csr.prmd := prmd.value()
    io.diff.csr.euen := euen.value()
    io.diff.csr.ecfg := ectl.value()
    io.diff.csr.estat := estat.value()
    io.diff.csr.era := era.value()
    io.diff.csr.badv := badv.value()
    io.diff.csr.eentry := eentry.value()
    io.diff.csr.tlbidx := tlbidx.value()
    io.diff.csr.tlbehi := tlbehi.value()
    io.diff.csr.tlbelo0 := tlbelo0.value()
    io.diff.csr.tlbelo1 := tlbelo1.value()
    io.diff.csr.asid := asid.value()
    io.diff.csr.pgdl := pgdl.value()
    io.diff.csr.pgdh := pgdh.value()
    io.diff.csr.save0 := save0.value()
    io.diff.csr.save1 := save1.value()
    io.diff.csr.save2 := save2.value()
    io.diff.csr.save3 := save3.value()
    io.diff.csr.tid := tid.value()
    io.diff.csr.tcfg := tcfg.value()
    io.diff.csr.tval := tval.value()
    io.diff.csr.ticlr := ticlr.value()
    io.diff.csr.llbctl := llbctl.value()
    io.diff.csr.tlbrentry := tlbrentry.value()
    io.diff.csr.dmw0 := dmw0.value()
    io.diff.csr.dmw1 := dmw1.value()

    io.diff.reg.zipWithIndex.foreach{ case (reg, i) => reg := GR(i)}
  }
}
